"""
    EnKFCore

Minimal, self-contained Ensemble Kalman Filter (EnKF) driver used by the
GSICoreAnalysis high-level orchestration. The implementation focuses on a
well-tested, deterministic square-root EnKF suitable for unit and integration
coverage while the full-featured driver is completed.

The driver accepts ensemble members supplied as either a flattened
`state_size × n_members` matrix or multidimensional grids with the ensemble
index on the final dimension. Observation dictionaries can either describe a
single consolidated observation set (`values`, `errors`, and `operator` or
`state_indices`) or a collection of observation groups with the same schema.

Returned analysis products are shaped to match the diagnostics module
requirements (`analysis_ensemble` as a 4D array, innovation vectors, and spread
metrics) so existing tooling continues to operate.
"""
module EnKFCore

using LinearAlgebra
using Statistics

import Random

export enkf_analysis, calculate_ensemble_covariances

# -----------------------------------------------------------------------------
# Public API
# -----------------------------------------------------------------------------

"""
    enkf_analysis(ensemble_data::Dict, observations::Dict, config)

Run a deterministic square-root Ensemble Kalman Filter update.

# Arguments
- `ensemble_data`: Dictionary containing ensemble members. A key named
  `"members"` is preferred; otherwise the first array-valued entry is used.
- `observations`: Observation dictionary describing values, errors, and either a
  linear observation operator matrix or state indices.
- `config`: Analysis configuration (only grid metadata and precision are used).

# Returns
A dictionary with the fields expected by the diagnostics layer:
- `"analysis_ensemble"`: 4D array (lon × lat × lev × member)
- `"analysis_mean"`: Analysis mean reshaped to the grid
- `"forecast_mean"`: Forecast mean reshaped to the grid
- `"innovation_vector"`: O−B departures
- `"innovation_fits"`: O−A departures
- `"ensemble_spread"`: Posterior spread in observation space
- `"observation_errors"`: Observation error standard deviations
- `"kalman_gain"`: Kalman gain matrix
- `"observation_operator"`: Linear observation operator used in the update
"""
function enkf_analysis(ensemble_data::Dict, observations::Dict, config)
    forecast_array = _locate_ensemble_array(ensemble_data)
    forecast_matrix, spatial_shape = _flatten_ensemble(forecast_array)

    n_state, n_members = size(forecast_matrix)
    n_members >= 2 || error("EnKF analysis requires at least two ensemble members")

    obs_values, obs_errors, H = _unpack_observations(observations, n_state)
    n_obs = length(obs_values)
    size(H, 2) == n_state || error("Observation operator has incompatible state dimension")
    length(obs_errors) == n_obs || error("Observation error vector length mismatch")

    # Prior moments
    ensemble_mean = vec(mean(forecast_matrix, dims = 2))
    anomalies = forecast_matrix .- ensemble_mean

    # Forecast covariance via ensemble statistics
    Pf = (anomalies * anomalies') / (n_members - 1)

    R = Diagonal(obs_errors .^ 2)
    HPf = H * Pf
    S = HPf * H' + R
    # Symmetrize and add a small ridge to protect against round-off issues
    S = (S + S') / 2
    ridge = 1e-9 * maximum(diag(S))
    if ridge > 0
        S += ridge * I
    end

    K = (Pf * H') * inv(S)

    innovation = obs_values - H * ensemble_mean
    analysis_mean = ensemble_mean + K * innovation

    # Deterministic square-root update for ensemble perturbations
    HA = H * anomalies
    analysis_anomalies = anomalies - K * HA
    analysis_matrix = analysis_mean .+ analysis_anomalies

    # Observation-space diagnostics
    analysis_obs = H * analysis_matrix
    analysis_obs_mean = vec(mean(analysis_obs, dims = 2))
    innovation_fits = obs_values - analysis_obs_mean
    ensemble_spread_obs = vec(std(analysis_obs, dims = 2; corrected = true))

    analysis_array = _reshape_to_4d(analysis_matrix, spatial_shape)
    forecast_array_4d = _reshape_to_4d(forecast_matrix, spatial_shape)
    analysis_mean_field = _reshape_mean(analysis_mean, spatial_shape)
    forecast_mean_field = _reshape_mean(ensemble_mean, spatial_shape)

    return Dict(
        "analysis_ensemble" => analysis_array,
        "analysis_mean" => analysis_mean_field,
        "forecast_mean" => forecast_mean_field,
        "innovation_vector" => innovation,
        "innovation_fits" => innovation_fits,
        "ensemble_spread" => ensemble_spread_obs,
        "observation_errors" => obs_errors,
        "kalman_gain" => K,
        "observation_operator" => H,
        "forecast_ensemble_matrix" => forecast_matrix,
        "analysis_ensemble_matrix" => analysis_matrix
    )
end

"""
    calculate_ensemble_covariances(ensemble_data::Dict, config)

Compute the ensemble mean and covariance matrix from the supplied members.
Returns a dictionary that can be consumed by the hybrid pathway while the
full operator-based implementation is under development.
"""
function calculate_ensemble_covariances(ensemble_data::Dict, config)
    ensemble_array = _locate_ensemble_array(ensemble_data)
    ensemble_matrix, _ = _flatten_ensemble(ensemble_array)

    n_state, n_members = size(ensemble_matrix)
    n_members >= 2 || error("At least two ensemble members are required")

    ensemble_mean = vec(mean(ensemble_matrix, dims = 2))
    anomalies = ensemble_matrix .- ensemble_mean
    covariance = (anomalies * anomalies') / (n_members - 1)

    precision_type = hasproperty(config, :precision) ? getfield(config, :precision) : Float64

    return Dict(
        "mean" => ensemble_mean,
        "covariance" => covariance,
        "anomalies" => anomalies,
        "precision" => precision_type
    )
end

# -----------------------------------------------------------------------------
# Helper utilities
# -----------------------------------------------------------------------------

function _locate_ensemble_array(ensemble_data::Dict)
    if haskey(ensemble_data, "members")
        arr = ensemble_data["members"]
        isa(arr, AbstractArray) || error("\"members\" entry must be an array")
        return arr
    end

    for value in values(ensemble_data)
        if isa(value, AbstractArray{<:Real})
            return value
        end
    end

    error("Ensemble data must contain a numeric array keyed by \"members\" or another variable entry")
end

function _flatten_ensemble(arr::AbstractArray)
    dims = size(arr)
    n_dims = ndims(arr)
    n_members = dims[end]
    state_dims = n_dims == 1 ? (dims[1],) : dims[1:end-1]
    state_size = prod(state_dims)
    matrix = reshape(arr, state_size, n_members)
    return matrix, state_dims
end

function _reshape_to_4d(matrix::AbstractMatrix, spatial_shape::Tuple)
    state_size, n_members = size(matrix)
    nx, ny, nz = _infer_spatial_dims(spatial_shape, state_size)
    return reshape(matrix, nx, ny, nz, n_members)
end

function _reshape_mean(vec::AbstractVector, spatial_shape::Tuple)
    state_size = length(vec)
    nx, ny, nz = _infer_spatial_dims(spatial_shape, state_size)
    return reshape(vec, nx, ny, nz)
end

function _infer_spatial_dims(spatial_shape::Tuple, state_size::Int)
    if length(spatial_shape) >= 3
        nx = spatial_shape[1]
        ny = spatial_shape[2]
        nz = prod(spatial_shape[3:end])
    elseif length(spatial_shape) == 2
        nx, ny = spatial_shape
        nz = cld(state_size, nx * ny)
    elseif length(spatial_shape) == 1
        nx = spatial_shape[1]
        ny = 1
        nz = cld(state_size, nx * ny)
    else
        nx = state_size
        ny = 1
        nz = 1
    end

    nx * ny * nz == state_size || error("Unable to infer spatial dimensions for ensemble reshaping")
    return nx, ny, nz
end

function _unpack_observations(observations::Dict, state_size::Int)
    if haskey(observations, "values")
        return _extract_observation_group(observations, state_size)
    end

    for value in values(observations)
        if isa(value, Dict) && haskey(value, "values")
            return _extract_observation_group(value, state_size)
        end
    end

    error("Observations dictionary must contain a group with \"values\" and \"errors\"")
end

function _extract_observation_group(group::Dict, state_size::Int)
    obs_values = Vector{Float64}(group["values"])
    obs_errors = haskey(group, "errors") ? Vector{Float64}(group["errors"]) : ones(Float64, length(obs_values))

    H = if haskey(group, "operator")
        Matrix{Float64}(group["operator"])
    elseif haskey(group, "state_indices")
        indices = collect(Int, group["state_indices"])
        n_obs = length(indices)
        mat = zeros(Float64, n_obs, state_size)
        for (row, col) in enumerate(indices)
            1 <= col <= state_size || error("State index $col out of bounds for observation operator")
            mat[row, col] = 1.0
        end
        mat
    else
        n_obs = length(obs_values)
        n_obs <= state_size || error("Cannot build identity observation operator with more observations than state elements")
        mat = zeros(Float64, n_obs, state_size)
        for i in 1:n_obs
            mat[i, i] = 1.0
        end
        mat
    end

    return obs_values, obs_errors, H
end

end # module EnKFCore
